#ifndef STX_ALGORITHM_NTH_ELEMENT_HPP
#define STX_ALGORITHM_NTH_ELEMENT_HPP

#include <stx/algorithm/common.hpp>
#include <stx/algorithm/swap.hpp>
#include <stx/algorithm/sort.hpp>

namespace stx {

template <class Compare, class RandomAccessIterator>
inline void
nth_element(
    RandomAccessIterator first,
    RandomAccessIterator nth,
    RandomAccessIterator last,
    Compare comp)
{
    // Compare is known to be a reference type
    typedef typename std::iterator_traits<RandomAccessIterator>::difference_type difference_type;
    const difference_type limit = 7;
    while (1) {
    restart_:
        difference_type len = last - first;
        switch (len) {
        case 0:
        case 1:
            return;
        case 2:
            if (comp(*--last, *first))
                swap(*first, *last);
            return;
        case 3:
            {
                RandomAccessIterator m = first;
                stx::sort3_impl<Compare>(first, ++m, --last, comp);
                return;
            }
        }
        if (len <= limit) {
            selection_sort_impl<Compare>(first, last, comp);
            return;
        }
        // len > limit >= 3
        RandomAccessIterator m = first + len/2;
        RandomAccessIterator lm1 = last;
        unsigned n_swaps = stx::sort3_impl<Compare>(first, m, --lm1, comp);
        // *m is median
        // partition [first, m) < *m and *m <= [m, last)
        // (this inhibits tossing elements equivalent to m around unnecessarily)
        RandomAccessIterator i = first;
        RandomAccessIterator j = lm1;
        // j points beyond range to be tested, *lm1 is known to be <= *m
        // The search going up is known to be guarded but the search coming down isn't.
        // Prime the downward search with a guard.
        if (!comp(*i, *m)) { // if *first == *m
            // *first == *m, *first doesn't go in first part
            // manually guard downward moving j against i
            while (1) {
                if (i == --j) {
                    // *first == *m, *m <= all other elements
                    // Parition instead into [first, i) == *first and *first < [i, last)
                    ++i;  // first + 1
                    j = last;
                    if (!comp(*first, *--j)) { // we need a guard if *first == *(last-1)
                        while (1) {
                            if (i == j) {
                                return;  // [first, last) all equivalent elements
                            }
                            if (comp(*first, *i)) {
                                swap(*i, *j);
                                ++n_swaps;
                                ++i;
                                break;
                            }
                            ++i;
                        }
                    }
                    // [first, i) == *first and *first < [j, last) and j == last - 1
                    if (i == j) {
                        return;
                    }
                    while (1) {
                        while (!comp(*first, *i)) {
                            ++i;
                        }
                        while (comp(*first, *--j)) {
                        }
                        if (i >= j) {
                            break;
                        }
                        swap(*i, *j);
                        ++n_swaps;
                        ++i;
                    }
                    // [first, i) == *first and *first < [i, last)
                    // The first part is sorted,
                    if (nth < i) {
                        return;
                    }
                    // nth_element the secod part
                    // nth_element<Compare>(i, nth, last, comp);
                    first = i;
                    goto restart_;
                }
                if (comp(*j, *m)) {
                    swap(*i, *j);
                    ++n_swaps;
                    break;  // found guard for downward moving j, now use unguarded partition
                }
            }
        }
        ++i;
        // j points beyond range to be tested, *lm1 is known to be <= *m
        // if not yet partitioned...
        if (i < j) {
            // known that *(i - 1) < *m
            while (1) {
                // m still guards upward moving i
                while (comp(*i, *m)) {
                    ++i;
                }
                // It is now known that a guard exists for downward moving j
                while (!comp(*--j, *m)) {
                }
                if (i >= j) {
                    break;
                }
                swap(*i, *j);
                ++n_swaps;
                // It is known that m != j
                // If m just moved, follow it
                if (m == i) {
                    m = j;
                }
                ++i;
            }
        }
        // [first, i) < *m and *m <= [i, last)
        if (i != m && comp(*m, *i)) {
            swap(*i, *m);
            ++n_swaps;
        }
        // [first, i) < *i and *i <= [i+1, last)
        if (nth == i) {
            return;
        }
        if (n_swaps == 0) {
            // We were given a perfectly partitioned sequence.  Coincidence?
            if (nth < i) {
                // Check for [first, i) already sorted
                j = m = first;
                while (++j != i) {
                    if (comp(*j, *m)) {
                        // not yet sorted, so sort
                        goto not_sorted;
                    }
                    m = j;
                }
                // [first, i) sorted
                return;
            } else {
                // Check for [i, last) already sorted
                j = m = i;
                while (++j != last) {
                    if (comp(*j, *m)) {
                        // not yet sorted, so sort
                        goto not_sorted;
                    }
                    m = j;
                }
                // [i, last) sorted
                return;
            }
        }
not_sorted:
        // nth_element on range containing nth
        if (nth < i) {
            // nth_element<Compare>(first, nth, i, comp);
            last = i;
        } else {
            // nth_element<Compare>(i+1, nth, last, comp);
            first = ++i;
        }
    }
}

template <class RandomAccessIterator, class Compare>
inline void
nth_element(
    RandomAccessIterator first,
    RandomAccessIterator nth,
    RandomAccessIterator last,
    Compare comp)
{
    typedef typename add_lvalue_reference<Compare>::type comp_ref;
    nth_element<comp_ref>(first, nth, last, comp);
}

template <class RandomAccessIterator>
inline void
nth_element(
    RandomAccessIterator first,
    RandomAccessIterator nth,
    RandomAccessIterator last)
{
    stx::nth_element(
        first,
        nth,
        last,
        stx::less<typename std::iterator_traits<RandomAccessIterator>::value_type>());
}

} // namespace stx

#endif // STX_ALGORITHM_NTH_ELEMENT_HPP
